#include <iostream>
#include <cstring>  // For memset
#include "test.pb.h"  // Protobuf 生成的头文件
#include <sodium.h>
#include <stdexcept>

// 固定大小的缓冲区大小
const size_t BUFFER_SIZE = 4096;

//// 加密函数
//void encryptWithLibsodium(const uint8_t* plaintext, size_t plaintext_size, unsigned char* key, unsigned char* nonce, unsigned char* ciphertext) {
//    // ChaCha20-Poly1305 加密
//    crypto_aead_chacha20poly1305_encrypt(ciphertext, nullptr, plaintext, plaintext_size, nullptr, 0, nullptr, nonce, key);
//}

void encryptWithLibsodium(const uint8_t* plaintext, size_t plaintext_size, unsigned char* key, unsigned char* nonce, unsigned char* ciphertext, size_t& ciphertext_size) {
    ciphertext_size = plaintext_size + crypto_aead_chacha20poly1305_ABYTES; // 计算加密后的数据长度
    if (crypto_aead_chacha20poly1305_encrypt(ciphertext, nullptr, plaintext, plaintext_size, nullptr, 0, nullptr, nonce, key) != 0) {
        throw std::logic_error("Encryption failed");
    }
}

// 解密函数
void decryptWithLibsodium(const unsigned char* ciphertext, size_t ciphertext_size, unsigned char* key, unsigned char* nonce, unsigned char* decrypted, unsigned long long& decrypted_size) {
    // 解密
    if (crypto_aead_chacha20poly1305_decrypt(decrypted, &decrypted_size, nullptr, ciphertext, ciphertext_size, nullptr, 0, nonce, key) != 0) {
        throw std::logic_error("Decryption failed");
    }
}

int main() {
    // 初始化 Libsodium
    if (sodium_init() < 0) {
        std::cerr << "Libsodium init failed" << std::endl;
        return -1;
    }

    // 创建 Protobuf 消息
    TestEncrypt message;
    message.set_content("加密解密成功");

    // 计算序列化后的消息大小
    int message_size = message.ByteSizeLong();

    // 检查消息是否适合固定大小的缓冲区
    if (message_size > BUFFER_SIZE) {
        std::cerr << "Message size exceeds buffer size" << std::endl;
        return -1;
    }

    // 使用固定大小的缓冲区
    uint8_t buf[BUFFER_SIZE];
    unsigned char key[crypto_aead_chacha20poly1305_KEYBYTES];
    unsigned char nonce[crypto_aead_chacha20poly1305_NPUBBYTES];
    unsigned char ciphertext[BUFFER_SIZE + crypto_aead_chacha20poly1305_ABYTES];
    size_t ciphertext_size;
    unsigned char decrypted[BUFFER_SIZE];
    unsigned long long decrypted_size;

    // 序列化消息到数组
    if (!message.SerializeToArray(buf, BUFFER_SIZE)) {
        std::cerr << "Failed to serialize Protobuf message" << std::endl;
        return -1;
    }

    // 生成密钥和 nonce
    randombytes_buf(key, sizeof key);
    randombytes_buf(nonce, sizeof nonce);


    // 计算1000万次加密耗时
    // 开始时间
    auto start = std::chrono::high_resolution_clock::now();
    for (int i = 0; i < 10000000; i++) {
        // 加密
        encryptWithLibsodium(buf, message_size, key, nonce, ciphertext, ciphertext_size);
    }
    auto end = std::chrono::high_resolution_clock::now();
    std::cout << "加密 1000万次 耗时: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;


    // 解密消息
    // size_t ciphertext_size = message_size + crypto_aead_chacha20poly1305_ABYTES;
    try {

        // 计算1000完此解密耗时
        start = std::chrono::high_resolution_clock::now();
        for (int i = 0; i < 10000000; i++) {
            // 解密
            decryptWithLibsodium(ciphertext, ciphertext_size, key, nonce, decrypted, decrypted_size);
        }
        end = std::chrono::high_resolution_clock::now();
        std::cout << "解密 1000万次 耗时: " << std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count() << "ms" << std::endl;

    } catch (const std::runtime_error& e) {
        std::cerr << e.what() << std::endl;
        return -1;
    }

    // 反序列化消息
    TestEncrypt decryptedMessage;
    if (!decryptedMessage.ParseFromArray(decrypted, decrypted_size)) {
        std::cerr << "Failed to parse Protobuf message" << std::endl;
        return -1;
    }

    // 打印解密后的内容
    std::cout << "解密出来的消息内容是: " << decryptedMessage.content() << std::endl;

    return 0;
}
